# email_filter.py
# 
# This file is the "bookkeeping" part of Peter Ballard's email spam filter.
# This is the top level program.
# It reads the user options, opens and closes the files, and sends emails
# to the appropriate files.
# However it is the function check() in the module email_spam, which this
# program calls, which does the actual checking.
#
# Feel free to modify this program to your taste.
# But the intention is that you can use this program unchanged,
# and that you only need to modify email_defs.py
#
# Version 1.01, 7-Apr-2005:
#   - all messages to stdout;
#   - neater way to clear srcfilename
# Version 1.0, 4-Apr-2005.
#
#
#Copyright (C) 2005  Peter Ballard
#
#This program is free software; you can redistribute it and/or
#modify it under the terms of the GNU General Public License
#as published by the Free Software Foundation; either version 2
#of the License, or (at your option) any later version.
#
#This program is distributed in the hope that it will be useful,
#but WITHOUT ANY WARRANTY; without even the implied warranty of
#MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#GNU General Public License for more details.
#
#You should have received a copy of the GNU General Public License
#along with this program; if not, write to the Free Software
#Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
# 

# Standard Python modules
import sys
import string
import re
import email
import email.Errors
#import mailbox
import os

# Python modules which are part of this package
import email_defs
import email_spam

##########################################
# Set up the options
##########################################

# set up the defaults (for non-test mode)
REAL = 1
srcfilename = email_defs.srcfilename
storefilename = email_defs.storefilename
spamfilename = email_defs.spamfilename
hamfilename = email_defs.hamfilename
whamfilename = email_defs.whamfilename
mhamfilename = email_defs.mhamfilename
panicfilename = email_defs.panicfilename
spammode = "ab"
hammode = "ab"
whammode = "ab"
mhammode = "ab"
panicmode = "ab"
detail = 0
summary = 1
verbose = 1
limit = 0

argv = sys.argv

if argv.count("-u"):
    sys.stdout.write("python ~/software/email_filter.py")
    sys.stdout.write(" [-test] \\\n")
    sys.stdout.write("          [-src srcfilename] (default is " + srcfilename + ") \\\n")
    sys.stdout.write("          [-[a]spam spamfilename] (default is " + spamfilename + ") \\\n")
    sys.stdout.write("          [-[a]ham hamfilename] (default is " + hamfilename + ") \\\n")
    sys.stdout.write("          [-[a]wham whamfilename] (default is " + whamfilename + ") \\\n")
    sys.stdout.write("          [-[a]mham whamfilename] (default is " + mhamfilename + ") \\\n")
    sys.stdout.write("          [-[a]panic panicfilename] (default is " + panicfilename + ") \\\n")
    sys.stdout.write("          [-limit n] (default is " + str(limit) + ") \\\n")
    sys.stdout.write("          [-detail n] (default is " + str(detail) + ") \\\n")
    sys.stdout.write("          [-summary n] (default is " + str(summary) + ") \\\n")
    sys.stdout.write("          [-verbose n] (default is " + str(verbose) + ")n")
    sys.exit(0)

if argv.count("-test"):
    # A different set of defaults for test mode
    REAL = 0
    argv.remove("-test")
    srcfilename = email_defs.test_srcfilename
    spamfilename = email_defs.test_spamfilename
    hamfilename = email_defs.test_hamfilename
    whamfilename = email_defs.test_whamfilename
    mhamfilename = email_defs.test_mhamfilename
    panicfilename = email_defs.test_panicfilename
    spammode = "wb"
    hammode = "wb"
    whammode = "wb"
    mhammode = "wb"
    panicmode = "wb"
    detail = 1
    summary = 1
    verbose = 1
    limit = 0

options = ["srcfilename", "spamfilename", "whamfilename", "mhamfilename",
           "hamfilename", "panicname", 
           "spammode", "whammode", "mhammode", "hammode", "panicmode",
           "detail", "summary", "verbose", "limit"]
    
if REAL==0:
    i = 1
    while i<len(argv):
        if i+1<len(argv) and argv[i][0]=="-" and (argv[i][1:] in options):
            # evaluate - no checking here!
            exec (argv[i][1:] + " = " + argv[i+1])
            i += 2
        else:
            raise Exception, "Bad input option " + argv[i]
elif len(argv)>1:
    raise Exception, "Option " + argv[1] + " only works with -test"

##########################################
# open the various files
##########################################

fp = open(srcfilename, "rb")
spamfile = open(spamfilename, spammode)
hamfile = open(hamfilename, hammode)
if whamfilename==hamfilename:
    whamfile = hamfile
else:
    whamfile = open(whamfilename, whammode)
if mhamfilename!=whamfilename:
    mhamfile = open(mhamfilename, mhammode)
    
panicfile = open(panicfilename, panicmode)
    
##########################################
# Read the source file (fp),
# Creating a list, msgs.
# Each list item is a string which is a single email message.
##########################################
msgs = [] # each list element is an email message 
thismsg = [] # each list element is a line
while 1:
    line = fp.readline()
    if line=="":
        break
    if len(line)>=5 and line[:5]=="From ":
        if len(thismsg):
            # this is MUCH faster than a loop of string concatenations,
            # because a big (immutable) string is only built once.
            msgs.append(string.join(thismsg, ""))
        thismsg = [line]
    else:
        thismsg.append(line)

# end condition
if len(thismsg):
    msgs.append(string.join(thismsg, ""))
fp.close()

# shorten the list if limit is specified.
# Perhaps this should be done dynamically while reading fp,
# So that msgs does not get ridiculously long.
if limit:
    msgs = msgs[-limit:]

##########################################
# Setup up the arrays WHITEHITS and BLACKHITS
# These are only used in the detailed summary if detail==1
##########################################

WHITEHITS = {}
BLACKHITS = {}

for word in email_spam.whitelist + email_spam.whiteintroregexps + email_spam.whiteintrostrings:
    WHITEHITS[word] = 0

for word in (email_spam.blackstrings + email_spam.blackregexps + email_spam.blackdomains
             + email_spam.black_attachment_types):
    BLACKHITS[word] = 0

##########################################
# process each email using results from the email_spam module.
# The algorithm is:
# if (message cannot be parsed):
#     it's a "panic"
# elif (sender is in whitelist):
#     it's "wham"
# elif (mail contains a string indicating it's good):
#     it's "mham"
# elif (mail looks like spam):
#     it's "spam"
# else:
#     it's "ham"
##########################################

spams = 0
hams = 0
whams = 0
mhams = 0
panics = 0
for msgstring in msgs:
    try:
        msg = email.message_from_string(msgstring)
    except:
        if verbose:
            sys.stdout.write("panic\n")
        panicfile.write(msgstring)
        panics += 1
        continue

    (status, matchword) = email_spam.check(msg)
    
    if status=="wham":
        whams = whams + 1
        if whamfile:
            whamfile.write(msgstring)
            #Do this to add a header...
            #msg.add_header("X-diyfilter", "wham")
            #whamfile.write(msg.as_string(1))
        if not WHITEHITS.has_key(matchword):
            WHITEHITS[matchword] = 1
        else:
            WHITEHITS[matchword] += 1
    elif status=="mham":
        mhams = mhams + 1
        # always write to whamfile, and optionally ALSO to mhamfile
        # (as a record of addresses that need to be whitelisted)
        whamfile.write(msgstring)
        if mhamfilename!=whamfilename:
            #Do this to add a header...
            #msg.add_header("X-diyfilter", "mham")
            #mhamfile.write(msg.as_string(1))
            mhamfile.write(msgstring)
        if not WHITEHITS.has_key(matchword):
            WHITEHITS[matchword] = 1
        else:
            WHITEHITS[matchword] += 1
    elif status=="spam":
        spams = spams+1
        #Do this to add a header...
        #msg.add_header("X-diyfilter", "spam")
        #spamfile.write(msg.as_string(1))
        spamfile.write(msgstring)
        if not BLACKHITS.has_key(matchword):
            BLACKHITS[matchword] = 1
        else:
            BLACKHITS[matchword] += 1
    else:
        #Do this to add a header...
        #msg.add_header("X-diyfilter", "ham")
        #hamfile.write(msg.as_string(1))
        hamfile.write(msgstring)
        hams = hams+1
        
    if verbose:
        # print a one-line summary per email
        sys.stdout.write(status + "\t" + matchword + "\t")
        sys.stdout.write(email_spam.subject(msg) + "\n")

if summary:
    sys.stdout.write(`len(msgs)` + " messages: "
                     + `whams` + " whams, " + `mhams` + " mhams, "
                     + `hams` + " hams, " + `spams` + " spams, " 
                     + `panics` + " panics\n")

spamfile.close()
hamfile.close()
panicfile.close()
if whamfilename!=hamfilename:
    whamfile.close()
if mhamfilename!=whamfilename:
    mhamfile.close()

##########################################
# For real email reading (as opposed to testing)
# empty srcfilename into storefilename.
##########################################

if REAL:
    os.system("cat " + srcfilename + " >> " + storefilename)
    # reduce srcfilename to zero size
    fp = open(srcfilename, "w")
    fp.close()
    # the old way
    #os.system("/bin/rm -f " + srcfilename)
    #os.system("touch " + srcfilename)

##########################################
# detailed summary, if requested
##########################################

if detail:
    for marker in ["+", "-"]:
        if marker=="+":
            dict = WHITEHITS
        else:
            dict = BLACKHITS
        tuples = []
        for key in dict.keys():
            tuples.append((dict[key], key))
        tuples.sort()
        tuples.reverse()
        for tuple in tuples:
            sys.stdout.write(marker + `tuple[0]` + "\t" + tuple[1] + "\n")
